import cv2 as cv
import numpy as np


# 内容：稠密光流估计
cap = cv.VideoCapture("/Users/apple/Desktop/data/vtest.avi")
ret, prev = cap.read()
prevGray = cv.cvtColor(prev, cv.COLOR_BGR2GRAY)


# 遍历当前帧每一行的第1列
hsv = np.zeros_like(prev)
hsv[..., 1] = 255


# 遍历视频每一帧
while ret:

    ret, next = cap.read()
    nextGray = cv.cvtColor(next, cv.COLOR_BGR2GRAY)

    # 把像素位移值，返回为两通道的光流向量
    flow = cv.calcOpticalFlowFarneback(prevGray, nextGray, None, 0.5, 3, 6, 3, 5, 1.25, 0)

    # 绘制光流：笛卡尔坐标转极坐标，得极轴和极角
    mag, ang = cv.cartToPolar(flow[..., 0], flow[..., 1])
    hsv[..., 0] = ang*180/np.pi/2
    hsv[..., 2] = cv.normalize(mag, None, 0, 255, cv.NORM_MINMAX)
    rgb = cv.cvtColor(hsv, cv.COLOR_HSV2BGR)

    # 显示
    cv.imshow("1 - Video Detection", next)
    cv.imshow("2 - Optical HSV", rgb)

    # 更新前后两帧图像
    prevGray = nextGray

    # 按esc退出
    if cv.waitKey(10) & 0xff == 27:
        break


cap.release()
cv.destroyAllWindows()
